// ROS /home/szd/catkin_yrobot/src/ROS_object_cluster1-main/src/object_cluster1.cpp
#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <geometry_msgs/PoseArray.h>
#include <visualization_msgs/MarkerArray.h>
#include "object_cluster1/ClusterArray.h"

#include <jsk_recognition_msgs/BoundingBox.h>
#include <jsk_recognition_msgs/BoundingBoxArray.h>

// PCL
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/filters/filter.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/filters/passthrough.h>
#include <pcl/filters/crop_box.h>
#include <pcl/segmentation/extract_clusters.h>
#include <pcl/common/common.h>
#include <pcl/common/centroid.h>

// using eigen lib
#include <Eigen/Dense>
#define NODE_VERSION "object_cluster1-2023-04-13-1-1"

//#define LOG

ros::Publisher cluster_array_pub_;
ros::Publisher cloud_filtered_pub_;
ros::Publisher marker_array_pub_;
ros::Publisher pub_bounding_boxs_;

bool print_fps_;
float x_axis_min_;
float x_axis_max_;
float y_axis_min_;
float y_axis_max_;
float z_axis_min_;
float z_axis_max_;
int cluster_size_min_;
int cluster_size_max_;
ros::Time timeGetPoints_;

const int region_max_ = 50; // Change this value to match how far you want to detect.
int regions_[100];

int frames; clock_t start_time; bool reset = true;//fps

float downsamplerSize_=0.1;  //单位米

void cal_transform(Eigen::Matrix4f &trans_matrix, float yaw, float pitch, float roll, float x, float y, float z)
{
     trans_matrix(0, 0) = cos(pitch)*cos(roll);
     trans_matrix(0, 1) = -1*cos(yaw)*sin(roll)+sin(yaw)*sin(pitch)*cos(roll);
     trans_matrix(0, 2) = sin(yaw)*sin(roll)+cos(yaw)*sin(pitch)*cos(roll);
     trans_matrix(1, 0) = cos(pitch)*sin(roll);
     trans_matrix(1, 1) = cos(yaw)*cos(roll)+sin(yaw)*sin(pitch)*sin(roll);
     trans_matrix(1, 2) = -1*sin(yaw)*cos(roll)+cos(yaw)*sin(pitch)*sin(roll);
     trans_matrix(2, 0) = -sin(pitch);
     trans_matrix(2, 1) = sin(yaw)*cos(pitch);
     trans_matrix(2, 2) = cos(yaw)*cos(pitch);
     trans_matrix(0, 3) = x;
     trans_matrix(1, 3) = y;
     trans_matrix(2, 3) = z;
     trans_matrix(3, 0) = 0;
     trans_matrix(3, 1) = 0;
     trans_matrix(3, 2) = 0;
     trans_matrix(3, 3) = 1;
}

void pointCloudCallback(const sensor_msgs::PointCloud2::ConstPtr& ros_pc2_in) {
   //if((ros::Time::now()-timeGetPoints_).toSec()<0.1)
   if((ros::Time::now()-timeGetPoints_).toSec()<0.2)
   {
        return;
   }

  if(print_fps_)if(reset){frames=0;start_time=clock();reset=false;}//fps

  /*** Convert ROS message to PCL ***/
  //pcl::PointCloud<pcl::PointXYZI>::Ptr pcl_pc_in(new pcl::PointCloud<pcl::PointXYZI>);
  //pcl::fromROSMsg(*ros_pc2_in, *pcl_pc_in);

  pcl::PointCloud<pcl::PointXYZ> pcl_from_ros;
  pcl::fromROSMsg(*ros_pc2_in, pcl_from_ros);
  //std::cout << "PointCloud initial has: " << pcl_from_ros.points.size ()  << " data points." << std::endl; //*

  //@坐标转换
  //Eigen::Matrix4f trans_matrix;
  //cal_transform(trans_matrix, 0, -0.00, 0, 0, 0, 0);
  //pcl::transformPointCloud(*pcl_pc_in, *pcl_pc_in, trans_matrix);

  /*** Remove ground and ceiling ***/
  pcl::IndicesPtr pc_indices(new std::vector<int>);
  // pcl::PassThrough<pcl::PointXYZI> pt;
  // pt.setInputCloud(pcl_pc_in);
  // pt.setFilterFieldName("z");
  // pt.setFilterLimits(z_axis_min_, z_axis_max_);
  // pt.filter(*pc_indices);
  //@yy 降采样
  // Create the filtering object: downsample the dataset using a leaf size of 1cm
  pcl::VoxelGrid<pcl::PointXYZ> vg;
  pcl::PointCloud<pcl::PointXYZ>::Ptr scan_ptr(new pcl::PointCloud<pcl::PointXYZ>(pcl_from_ros));
  pcl::PointCloud<pcl::PointXYZ>::Ptr pcl_pc_in (new pcl::PointCloud<pcl::PointXYZ>);

  vg.setInputCloud (scan_ptr);

  vg.setLeafSize (downsamplerSize_, downsamplerSize_, downsamplerSize_);
  vg.filter (*pcl_pc_in);
  //std::cout << "PointCloud after filtering has: " << pcl_pc_in->points.size ()  << " data points." << std::endl; //*
  //simpleVis(cloud_filtered);



  //pcl::CropBox<pcl::PointXYZI> region;
  pcl::CropBox<pcl::PointXYZ> region;
  region.setMin(Eigen::Vector4f (x_axis_min_, y_axis_min_, z_axis_min_, 1.0));
  region.setMax(Eigen::Vector4f (x_axis_max_, y_axis_max_, z_axis_max_, 1.0));
  region.setInputCloud(pcl_pc_in);
  region.filter(*pc_indices);

  /*** Divide the point cloud into nested circular regions ***/
  boost::array<std::vector<int>, region_max_> indices_array;
  //printf("pc_indices->size()=%d\n",pc_indices->size());

  for(int i = 0; i < pc_indices->size(); i++) {
    float range = 0.0;
    for(int j = 0; j < region_max_; j++) {
      float d2 = pcl_pc_in->points[(*pc_indices)[i]].x * pcl_pc_in->points[(*pc_indices)[i]].x +
                       pcl_pc_in->points[(*pc_indices)[i]].y * pcl_pc_in->points[(*pc_indices)[i]].y +
                       pcl_pc_in->points[(*pc_indices)[i]].z * pcl_pc_in->points[(*pc_indices)[i]].z;
      if(d2 > range * range && d2 <= (range+regions_[j]) * (range+regions_[j])) {
        indices_array[j].push_back((*pc_indices)[i]);
        break;
      }
      range += regions_[j];
    }
  }
  /*** Euclidean clustering ***/
  float tolerance = 0.2;
  std::vector<pcl::PointCloud<pcl::PointXYZ>::Ptr, Eigen::aligned_allocator<pcl::PointCloud<pcl::PointXYZI>::Ptr > > clusters;

  for(int i = 0; i < region_max_; i++) {
    tolerance += 0.1;
    if(indices_array[i].size() > cluster_size_min_) {

      boost::shared_ptr<std::vector<int> > indices_array_ptr(new std::vector<int>(indices_array[i]));
      //pcl::search::KdTree<pcl::PointXYZ>::Ptr tree(new pcl::search::KdTree<pcl::PointXYZ>);
      pcl::search::KdTree<pcl::PointXYZ>::Ptr tree(new pcl::search::KdTree<pcl::PointXYZ>);

      tree->setInputCloud(pcl_pc_in, indices_array_ptr);

      std::vector<pcl::PointIndices> cluster_indices;
      pcl::EuclideanClusterExtraction<pcl::PointXYZ> ec;
      ec.setClusterTolerance(tolerance);
      ec.setMinClusterSize(cluster_size_min_);
      ec.setMaxClusterSize(cluster_size_max_);
      ec.setSearchMethod(tree);
      ec.setInputCloud(pcl_pc_in);
      ec.setIndices(indices_array_ptr);
      ec.extract(cluster_indices);

      for(std::vector<pcl::PointIndices>::const_iterator it = cluster_indices.begin(); it != cluster_indices.end(); it++) {

        pcl::PointCloud<pcl::PointXYZ>::Ptr cluster(new pcl::PointCloud<pcl::PointXYZ>);

        for(std::vector<int>::const_iterator pit = it->indices.begin(); pit != it->indices.end(); ++pit) {
          cluster->points.push_back(pcl_pc_in->points[*pit]);
            }

        cluster->width = cluster->size();
        cluster->height = 1;
        cluster->is_dense = true;
              clusters.push_back(cluster);
      }
    }
  }

  /*** Output ***/
  if(cloud_filtered_pub_.getNumSubscribers() > 0) {

    pcl::PointCloud<pcl::PointXYZ>::Ptr pcl_pc_out(new pcl::PointCloud<pcl::PointXYZ>);
    sensor_msgs::PointCloud2 ros_pc2_out;
    pcl::copyPointCloud(*pcl_pc_in, *pc_indices, *pcl_pc_out);
    pcl::toROSMsg(*pcl_pc_out, ros_pc2_out);
    cloud_filtered_pub_.publish(ros_pc2_out);

  }

  object_cluster1::ClusterArray cluster_array;
  geometry_msgs::PoseArray pose_array;
  visualization_msgs::MarkerArray marker_array;

  jsk_recognition_msgs::BoundingBoxArray bbox_array;
  std_msgs::Header point_cloud_header_;
  point_cloud_header_ = ros_pc2_in->header;

  for(int i = 0; i < clusters.size(); i++) {
    if(cluster_array_pub_.getNumSubscribers() > 0) {
      sensor_msgs::PointCloud2 ros_pc2_out;
      pcl::toROSMsg(*clusters[i], ros_pc2_out);
      cluster_array.clusters.push_back(ros_pc2_out);
    }

    if(marker_array_pub_.getNumSubscribers() > 0) {
      Eigen::Vector4f min, max;
      pcl::getMinMax3D(*clusters[i], min, max);

      visualization_msgs::Marker marker;
      marker.header = ros_pc2_in->header;
      marker.ns = "object_cluster1";
      marker.id = i;
      marker.type = visualization_msgs::Marker::LINE_LIST;

      geometry_msgs::Point p[24];
      p[0].x = max[0];  p[0].y = max[1];  p[0].z = max[2];
      p[1].x = min[0];  p[1].y = max[1];  p[1].z = max[2];
      p[2].x = max[0];  p[2].y = max[1];  p[2].z = max[2];
      p[3].x = max[0];  p[3].y = min[1];  p[3].z = max[2];
      p[4].x = max[0];  p[4].y = max[1];  p[4].z = max[2];
      p[5].x = max[0];  p[5].y = max[1];  p[5].z = min[2];
      p[6].x = min[0];  p[6].y = min[1];  p[6].z = min[2];
      p[7].x = max[0];  p[7].y = min[1];  p[7].z = min[2];
      p[8].x = min[0];  p[8].y = min[1];  p[8].z = min[2];
      p[9].x = min[0];  p[9].y = max[1];  p[9].z = min[2];
      p[10].x = min[0]; p[10].y = min[1]; p[10].z = min[2];
      p[11].x = min[0]; p[11].y = min[1]; p[11].z = max[2];
      p[12].x = min[0]; p[12].y = max[1]; p[12].z = max[2];
      p[13].x = min[0]; p[13].y = max[1]; p[13].z = min[2];
      p[14].x = min[0]; p[14].y = max[1]; p[14].z = max[2];
      p[15].x = min[0]; p[15].y = min[1]; p[15].z = max[2];
      p[16].x = max[0]; p[16].y = min[1]; p[16].z = max[2];
      p[17].x = max[0]; p[17].y = min[1]; p[17].z = min[2];
      p[18].x = max[0]; p[18].y = min[1]; p[18].z = max[2];
      p[19].x = min[0]; p[19].y = min[1]; p[19].z = max[2];
      p[20].x = max[0]; p[20].y = max[1]; p[20].z = min[2];
      p[21].x = min[0]; p[21].y = max[1]; p[21].z = min[2];
      p[22].x = max[0]; p[22].y = max[1]; p[22].z = min[2];
      p[23].x = max[0]; p[23].y = min[1]; p[23].z = min[2];
      for(int i = 0; i < 24; i++) {
            marker.points.push_back(p[i]);
      }

      //marker.scale.x = 0.02;
      marker.scale.x = 0.1;
      marker.color.a = 1.0;
      marker.color.r = 0.0;
      marker.color.g = 1.0;
      marker.color.b = 0.5;
      marker.lifetime = ros::Duration(0.1);
      marker_array.markers.push_back(marker);
    }

    if(pub_bounding_boxs_.getNumSubscribers() > 0) {
      Eigen::Vector4f min, max, centroid;
      pcl::getMinMax3D(*clusters[i], min, max);
      pcl::compute3DCentroid(*clusters[i], centroid);

      float minmax_x, minmax_y, minmax_z;
      minmax_x = (max[0] + min[0])/2.0;
      minmax_y = (max[1] + min[1])/2.0;
      minmax_z = (max[2] + min[2])/2.0;

      // float minmax_dist, centroid_dist;
      // minmax_dist = minmax_x*minmax_x + minmax_y*minmax_y;
      // centroid_dist = centroid[0]*centroid[0] + centroid[1]*centroid[1];

      // if(centroid_dist == minmax_dist){

      //   jsk_recognition_msgs::BoundingBox bounding_box_;

      //   bounding_box_.header = point_cloud_header_;

      //   bounding_box_.pose.position.x = minmax_x;
      //   bounding_box_.pose.position.y = minmax_y;
      //   bounding_box_.pose.position.z = minmax_z;
      //   bounding_box_.dimensions.x = max[0] - min[0];
      //   bounding_box_.dimensions.y = max[1] - min[1];
      //   bounding_box_.dimensions.z = max[2] - min[2];

      //   bbox_array.boxes.push_back(bounding_box_);
      // }
      // else{
      //   jsk_recognition_msgs::BoundingBox bounding_box_1_;
      //   jsk_recognition_msgs::BoundingBox bounding_box_2_;
      //   bounding_box_1_.header = point_cloud_header_;
      //   bounding_box_2_.header = point_cloud_header_;

      //   if((minmax_x-centroid[0])>=0){
      //     bounding_box_1_.pose.position.x = min[0];
      //     bounding_box_1_.pose.position.y = minmax_y;
      //     bounding_box_1_.pose.position.z = minmax_z;
      //     bounding_box_1_.dimensions.x = 0.2;
      //     bounding_box_1_.dimensions.y = max[1] - min[1];
      //     bounding_box_1_.dimensions.z = max[2] - min[2];
      //   }
      //   else{
      //     bounding_box_1_.pose.position.x = max[0];
      //     bounding_box_1_.pose.position.y = minmax_y;
      //     bounding_box_1_.pose.position.z = minmax_z;
      //     bounding_box_1_.dimensions.x = 0.2;
      //     bounding_box_1_.dimensions.y = max[1] - min[1];
      //     bounding_box_1_.dimensions.z = max[2] - min[2];
      //   }

      //   if((minmax_y-centroid[1])>=0){
      //     bounding_box_2_.pose.position.x = minmax_x;
      //     bounding_box_2_.pose.position.y = min[1];
      //     bounding_box_2_.pose.position.z = minmax_z;
      //     bounding_box_2_.dimensions.x = max[0] - min[0];
      //     bounding_box_2_.dimensions.y = 0.2;
      //     bounding_box_2_.dimensions.z = max[2] - min[2];
      //   }
      //   else{
      //     bounding_box_2_.pose.position.x = minmax_x;
      //     bounding_box_2_.pose.position.y = max[1];
      //     bounding_box_2_.pose.position.z = minmax_z;
      //     bounding_box_2_.dimensions.x = max[0] - min[0];
      //     bounding_box_2_.dimensions.y = 0.2;
      //     bounding_box_2_.dimensions.z = max[2] - min[2];
      //   }
      //   bbox_array.boxes.push_back(bounding_box_1_);
      //   bbox_array.boxes.push_back(bounding_box_2_);
      // }


      jsk_recognition_msgs::BoundingBox bounding_box_1_;
      jsk_recognition_msgs::BoundingBox bounding_box_2_;
      bounding_box_1_.header = point_cloud_header_;
      bounding_box_2_.header = point_cloud_header_;

      if((minmax_x-centroid[0])>=0){
        bounding_box_1_.pose.position.x = min[0];
        bounding_box_1_.pose.position.y = minmax_y;
        bounding_box_1_.pose.position.z = minmax_z;
        bounding_box_1_.dimensions.x = 0.2;
        bounding_box_1_.dimensions.y = max[1] - min[1];
        bounding_box_1_.dimensions.z = max[2] - min[2];
      }
      else{
        bounding_box_1_.pose.position.x = max[0];
        bounding_box_1_.pose.position.y = minmax_y;
        bounding_box_1_.pose.position.z = minmax_z;
        bounding_box_1_.dimensions.x = 0.2;
        bounding_box_1_.dimensions.y = max[1] - min[1];
        bounding_box_1_.dimensions.z = max[2] - min[2];
      }

      if((minmax_y-centroid[1])>=0){
        bounding_box_2_.pose.position.x = minmax_x;
        bounding_box_2_.pose.position.y = min[1];
        bounding_box_2_.pose.position.z = minmax_z;
        bounding_box_2_.dimensions.x = max[0] - min[0];
        bounding_box_2_.dimensions.y = 0.2;
        bounding_box_2_.dimensions.z = max[2] - min[2];
      }
      else{
        bounding_box_2_.pose.position.x = minmax_x;
        bounding_box_2_.pose.position.y = max[1];
        bounding_box_2_.pose.position.z = minmax_z;
        bounding_box_2_.dimensions.x = max[0] - min[0];
        bounding_box_2_.dimensions.y = 0.2;
        bounding_box_2_.dimensions.z = max[2] - min[2];
      }
      bbox_array.boxes.push_back(bounding_box_1_);
      bbox_array.boxes.push_back(bounding_box_2_);
    }
     timeGetPoints_=ros::Time::now();
  }

  if(cluster_array.clusters.size()) {
    cluster_array.header = ros_pc2_in->header;
    cluster_array_pub_.publish(cluster_array);
  }

  if(marker_array.markers.size()) {
    marker_array_pub_.publish(marker_array);
  }

  if(bbox_array.boxes.size()) {
    bbox_array.header = point_cloud_header_;
    printf("@box=%d\n",bbox_array.boxes.size());
    pub_bounding_boxs_.publish(bbox_array);
  }

  //if(print_fps_)if(++frames>10){std::cerr<<"[object_cluster1] fps = "<<float(frames)/(float(clock()-start_time)/CLOCKS_PER_SEC)<<", timestamp = "<<clock()/CLOCKS_PER_SEC<<std::endl;reset = true;}//fps
}

int main(int argc, char **argv) {
  ros::init(argc, argv, "object_cluster1");
  printf("###########################version###########################\n");
  std::cout<<NODE_VERSION<<std::endl;
  printf("###########################version###########################\n");


  /*** Subscribers ***/
  ros::NodeHandle nh;
  //ros::Subscriber point_cloud_sub = nh.subscribe<sensor_msgs::PointCloud2>("velodyne_points", 1, pointCloudCallback);
  ros::Subscriber point_cloud_sub = nh.subscribe<sensor_msgs::PointCloud2>("/os1_cloud_node/points", 1, pointCloudCallback);

  /*** Publishers ***/
  ros::NodeHandle private_nh("~");
  cluster_array_pub_ = private_nh.advertise<object_cluster1::ClusterArray>("clusters", 100);
  cloud_filtered_pub_ = private_nh.advertise<sensor_msgs::PointCloud2>("cloud_filtered", 100);
  marker_array_pub_ = private_nh.advertise<visualization_msgs::MarkerArray>("markers", 100);
  pub_bounding_boxs_ = private_nh.advertise<jsk_recognition_msgs::BoundingBoxArray>("detected_bounding_boxs", 5);

  /*** Parameters ***/
  std::string sensor_model;

  private_nh.param<std::string>("sensor_model", sensor_model, "64"); // VLP-16, HDL-32E, HDL-64E
  private_nh.param<bool>("print_fps", print_fps_, false);
  private_nh.param<float>("x_axis_min", x_axis_min_, 0.0);
  private_nh.param<float>("x_axis_max", x_axis_max_, 80.0);
  private_nh.param<float>("y_axis_min", y_axis_min_, -20.0);
  private_nh.param<float>("y_axis_max", y_axis_max_, 20.0);
  private_nh.param<float>("z_axis_min", z_axis_min_, -0.70);
  private_nh.param<float>("z_axis_max", z_axis_max_, 0.5);
  private_nh.param<int>("cluster_size_min", cluster_size_min_, 10); //50
  private_nh.param<int>("cluster_size_max", cluster_size_max_, 10000); //2200000

  // Divide the point cloud into nested circular regions centred at the sensor.
  // For more details, see our IROS-17 paper "Online learning for human classification in 3D LiDAR-based tracking"
  if(sensor_model.compare("16") == 0) {
    regions_[0] = 2; regions_[1] = 3; regions_[2] = 3; regions_[3] = 3; regions_[4] = 3;
    regions_[5] = 3; regions_[6] = 3; regions_[7] = 2; regions_[8] = 3; regions_[9] = 3;
    regions_[10]= 3; regions_[11]= 3; regions_[12]= 3; regions_[13]= 3;
  } else if (sensor_model.compare("32") == 0) {
    regions_[0] = 4; regions_[1] = 5; regions_[2] = 4; regions_[3] = 5; regions_[4] = 4;
    regions_[5] = 5; regions_[6] = 5; regions_[7] = 4; regions_[8] = 5; regions_[9] = 4;
    regions_[10]= 5; regions_[11]= 5; regions_[12]= 4; regions_[13]= 5;
  } else if (sensor_model.compare("64") == 0) {
    regions_[0] = 14; regions_[1] = 14; regions_[2] = 14; regions_[3] = 15; regions_[4] = 14;
  } else {
    ROS_FATAL("Unknown sensor model!");
  }

  ros::spin();

  return 0;
}
